import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import cv2
import json

from tqdm import tqdm
from typing import Dict, List, Tuple
from configs.config import load_config


class VideoProcessor:
    _image_cache: Dict[str, List[Tuple[int, str]]] = {}

    def create_combined_video(self, data_root_list: Dict[str, str], output_path: str, fps: int = 30) -> None:
        for data_name, root_path in tqdm(data_root_list.items(), desc="处理目录"):
            camera_dir = os.path.join(root_path, 'Camera')
            if not os.path.exists(camera_dir):
                print(f"警告: {camera_dir} 不存在")
                continue
                
            images = sorted(
                [f for f in os.listdir(camera_dir) if f.endswith('.jpg')],
                key=lambda x: int(x.split('_')[0])
            )
            
            if not images:
                print(f"错误: {camera_dir} 中没有找到图片")
                continue
                
            first_frame = cv2.imread(os.path.join(camera_dir, images[0]))
            if first_frame is None:
                print(f"错误: 无法读取第一张图片 {images[0]}")
                continue
                
            height, width, _ = first_frame.shape
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            video_path = os.path.join(output_path, f"{data_name}.mp4")
            os.makedirs(output_path, exist_ok=True)
            
            video_writer = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
            if not video_writer.isOpened():
                print(f"错误: 无法创建视频文件 {video_path}")
                continue
                
            for img in tqdm(images, desc=f"生成 {data_name} 视频"):
                img_path = os.path.join(camera_dir, img)
                frame = cv2.imread(img_path)
                if frame is not None:
                    video_writer.write(frame)
                else:
                    print(f"警告: 无法读取图片 {img_path}")
                    
            video_writer.release()
            print(f"视频生成完成: {video_path}")
    
    def create_event_time_markers(self, json_root_dir: str, data_root_list: Dict[str, str]) -> None:
        subdirs = [d for d in os.listdir(json_root_dir) if os.path.isdir(os.path.join(json_root_dir, d))]
        
        for subdir in tqdm(subdirs, desc="处理事件目录", unit="目录"):
            subdir_path = os.path.join(json_root_dir, subdir)
            
            data_name = subdir.split('_')[0]
            if data_name not in data_root_list:
                print(f"警告: {data_name} 没有匹配的数据根目录")
                continue
            
            image_dir = os.path.join(data_root_list[data_name], 'Camera')
            if image_dir not in VideoProcessor._image_cache:
                VideoProcessor._load_image_cache(image_dir)
            
            for file in os.listdir(subdir_path):
                if file.startswith("total_BEV_") and file.endswith(".json"):
                    json_path = os.path.join(subdir_path, file)
                    try:
                        with open(json_path, 'r', encoding='utf-8') as f:
                            data = json.load(f)
                        
                        start_timestamp = int(data["start_id"].split('_')[1])
                        image_list = VideoProcessor._image_cache[image_dir]
                        
                        if not image_list:
                            print(f'错误: {image_dir} 无图片缓存')
                            continue
                            
                        target_idx1 = -1
                        min_time_diff1 = float('inf')
                        for idx, (timestamp, _) in enumerate(image_list):
                            time_diff = abs(timestamp - start_timestamp)
                            if time_diff < min_time_diff1:
                                min_time_diff1 = time_diff
                                target_idx1 = idx
                        
                        target_idx2 = -1
                        min_time_diff2 = float('inf')
                        for idx, (timestamp, _) in enumerate(image_list):
                            time_diff = abs(timestamp - start_timestamp - 28800000)
                            if time_diff < min_time_diff2:
                                min_time_diff2 = time_diff
                                target_idx2 = idx

                        if min_time_diff1 >= 100000 and min_time_diff2 >= 100000:
                            raise ValueError(f"错误: {json_path} 中没有找到匹配的图片时间戳")
                        elif min_time_diff1 < 100000 and min_time_diff2 < 100000:
                            raise ValueError(f"错误: {json_path} 中找到多个匹配的图片时间戳")
                        elif min_time_diff1 < 100000:
                            target_idx = target_idx1
                        else:
                            target_idx = target_idx2

                        total_seconds = target_idx / 30.0
                        minutes = int(total_seconds // 60)
                        seconds = int(total_seconds % 60)
                        
                        file_name = f"{minutes:02d}_{seconds:02d}"
                        open(os.path.join(subdir_path, file_name), 'w').close()
                        
                    except Exception as e:
                        print(f"处理 {json_path} 错误: {str(e)}")
                    break

    @staticmethod
    def _load_image_cache(image_dir: str) -> None:
        image_list = []
        for img in os.listdir(image_dir):
            if not img.endswith('.jpg'):
                continue
            try:
                parts = img.split('_')
                timestamp = int(parts[1].split('.')[0])
                image_list.append((timestamp, img))
            except (IndexError, ValueError):
                continue
        image_list.sort()
        VideoProcessor._image_cache[image_dir] = image_list

if __name__ == "__main__":
    cfg = load_config("./configs/data_process.py")
    data_root_list = cfg.get('data_root_list')
    output_video_path = cfg.get('bev_output_folder')
    
    processor = VideoProcessor()
    processor.create_combined_video(data_root_list, output_video_path)
    processor.create_event_time_markers(output_video_path, data_root_list)